#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import sys
import numpy as np

# 对于float32的验证容差
relative_tol = 1e-4
absolute_tol = 1e-5
error_tol = 1e-4


def verify_result(output_mean, output_var, golden_mean, golden_var):
    # 读取输出结果和黄金标准
    out_mean = np.fromfile(output_mean, dtype=np.float32).reshape(-1)
    out_var = np.fromfile(output_var, dtype=np.float32).reshape(-1)
    gold_mean = np.fromfile(golden_mean, dtype=np.float32).reshape(-1)
    gold_var = np.fromfile(golden_var, dtype=np.float32).reshape(-1)
    
    # 检查形状是否一致
    if out_mean.shape != gold_mean.shape or out_var.shape != gold_var.shape:
        print(f"Shape mismatch: output mean {out_mean.shape} vs golden mean {gold_mean.shape}")
        print(f"Shape mismatch: output var {out_var.shape} vs golden var {gold_var.shape}")
        return False
    
    # 验证均值
    mean_close = np.isclose(out_mean, gold_mean, rtol=relative_tol, atol=absolute_tol)
    mean_error_ratio = 1.0 - np.mean(mean_close)
    
    # 验证方差
    var_close = np.isclose(out_var, gold_var, rtol=relative_tol, atol=absolute_tol)
    var_error_ratio = 1.0 - np.mean(var_close)
    
    # 打印错误信息
    if mean_error_ratio > 0:
        print(f"Mean verification failed. Error ratio: {mean_error_ratio:.6f}")
        bad_indices = np.where(~mean_close)[0][:10]  # 只显示前10个错误
        for idx in bad_indices:
            print(f"Index {idx}: expected {gold_mean[idx]:.6f}, got {out_mean[idx]:.6f}")
    
    if var_error_ratio > 0:
        print(f"Variance verification failed. Error ratio: {var_error_ratio:.6f}")
        bad_indices = np.where(~var_close)[0][:10]  # 只显示前10个错误
        for idx in bad_indices:
            print(f"Index {idx}: expected {gold_var[idx]:.6f}, got {out_var[idx]:.6f}")
    
    # 检查是否在容差范围内
    return mean_error_ratio <= error_tol and var_error_ratio <= error_tol

if __name__ == '__main__':
    try:
        if len(sys.argv) != 5:
            print("Usage: python3 verify_result.py output_mean.bin output_var.bin golden_mean.bin golden_var.bin")
            sys.exit(1)
            
        res = verify_result(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])
        if not res:
            raise ValueError("[ERROR] Result verification failed")
        else:
            print("Test pass")
    except Exception as e:
        print(e)
        sys.exit(1)
